{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# How to use Pretrained models with PyTorch\n",
    "# Simple Classifier using resnet50\n",
    "# code by GunhoChoi\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "import torchvision.models as models\n",
    "import torchvision.datasets as dset\n",
    "import torchvision.transforms as transforms\n",
    "from torch.autograd import Variable\n",
    "\n",
    "batch_size = 3\n",
    "learning_rate =0.0002\n",
    "epoch = 50\n",
    "\n",
    "resnet = models.resnet50(pretrained=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['dogs', 'onepun']\n",
      "{'onepun': 1, 'dogs': 0}\n",
      "[('./images/dogs/3b1a8d3fca2adb68a7a908b73005764a.jpg', 0), ('./images/dogs/dog-650299.jpg', 0), ('./images/onepun/5320.jpg', 1), ('./images/onepun/5720.jpg', 1), ('./images/onepun/1454.jpg', 1)]\n"
     ]
    }
   ],
   "source": [
    "# Input pipeline from a folder containing multiple folders of images\n",
    "# we can check the classes, class_to_idx, and filename with idx\n",
    "\n",
    "img_dir = \"./images\"\n",
    "img_data = dset.ImageFolder(img_dir, transforms.Compose([\n",
    "            transforms.RandomSizedCrop(224),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "            ]))\n",
    "\n",
    "print(img_data.classes)\n",
    "print(img_data.class_to_idx)\n",
    "print(img_data.imgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 3, 224, 224])\n",
      "\n",
      " 1\n",
      " 0\n",
      " 0\n",
      "[torch.LongTensor of size 3]\n",
      "\n",
      "torch.Size([2, 3, 224, 224])\n",
      "\n",
      " 1\n",
      " 1\n",
      "[torch.LongTensor of size 2]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# After we get the list of images, we can turn the list into batches of images\n",
    "# with torch.utils.data.DataLoader()\n",
    "\n",
    "img_batch = data.DataLoader(img_data, batch_size=batch_size,\n",
    "                            shuffle=True, num_workers=2)\n",
    "\n",
    "for img,label in img_batch:\n",
    "    print(img.size())\n",
    "    print(label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      "-8.9971e-01  1.8495e+00 -7.6565e-01  ...  -2.7083e-01  1.1463e+00  7.6773e-01\n",
      "-1.1474e+00 -2.7008e+00 -1.6100e-01  ...  -7.4485e-01  1.0327e+00  9.7581e-01\n",
      "[torch.FloatTensor of size 2x1000]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# test of the result coming from resnet model\n",
    "\n",
    "img = Variable(img)\n",
    "print(resnet(img))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      " 0.0000  0.0000\n",
      " 0.0202  0.1174\n",
      "[torch.FloatTensor of size 2x2]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# we have 2 categorical variables so 1000 -> 500 -> 2\n",
    "# test the whole process\n",
    "\n",
    "model = nn.Sequential(\n",
    "            nn.Linear(1000,500),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(500),\n",
    "            nn.Linear(500,2),\n",
    "            nn.ReLU()\n",
    "            )\n",
    "\n",
    "print(model(resnet(img)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# define loss func & optimizer\n",
    "\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(),lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Variable containing:\n",
      " 0.6731\n",
      "[torch.cuda.FloatTensor of size 1 (GPU 0)]\n",
      "\n",
      "Variable containing:\n",
      " 0.1374\n",
      "[torch.cuda.FloatTensor of size 1 (GPU 0)]\n",
      "\n",
      "Variable containing:\n",
      "1.00000e-02 *\n",
      "  8.1255\n",
      "[torch.cuda.FloatTensor of size 1 (GPU 0)]\n",
      "\n",
      "Variable containing:\n",
      "1.00000e-02 *\n",
      "  4.8874\n",
      "[torch.cuda.FloatTensor of size 1 (GPU 0)]\n",
      "\n",
      "Variable containing:\n",
      "1.00000e-02 *\n",
      "  4.2472\n",
      "[torch.cuda.FloatTensor of size 1 (GPU 0)]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# In order to train with GPU, we need to put the model and variables\n",
    "# by doing .cuda()\n",
    "\n",
    "resnet.cuda()\n",
    "model.cuda()\n",
    "\n",
    "for i in range(epoch):\n",
    "    for img,label in img_batch:\n",
    "        img = Variable(img).cuda()\n",
    "        label = Variable(label).cuda()\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        output = model(resnet(img))\n",
    "        loss = loss_func(output,label)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    if i % 10 ==0:\n",
    "        print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 1.0\n"
     ]
    }
   ],
   "source": [
    "# Check Accuracy of the trained model\n",
    "# Need to get used to using .cuda() and .data \n",
    "\n",
    "model.eval()\n",
    "correct = 0\n",
    "total = 0\n",
    "\n",
    "for img,label in img_batch:\n",
    "    img = Variable(img).cuda()\n",
    "    label = Variable(label).cuda()\n",
    "    \n",
    "    output = model(resnet(img))\n",
    "    _, pred = torch.max(output.data,1)\n",
    "    \n",
    "    total += label.size(0)\n",
    "    correct += (pred == label.data).sum()   \n",
    "\n",
    "print(\"Accuracy: {}\".format(correct/total))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
